#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
初始化任务场景数据脚本
"""
import sys
import os
import json

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database.db import SessionLocal, engine, Base
from models.task_scenario import TaskScenario
from models.detection_target import DetectionTarget

def init_task_scenarios():
    """初始化任务场景数据"""
    # 创建数据库表
    Base.metadata.create_all(bind=engine)
    
    # 创建数据库会话
    db = SessionLocal()
    
    try:
        # 获取所有检测目标，创建名称到ID的映射
        detection_targets = db.query(DetectionTarget).all()
        target_name_to_id = {target.name: target.id for target in detection_targets}
        
        # 检查是否已有预设场景
        existing_presets = db.query(TaskScenario).filter(
            TaskScenario.scenario_type == "preset"
        ).all()
        
        is_update_mode = len(existing_presets) > 0
        
        # 验证所有需要的检测目标是否存在
        required_targets = [
            "线路断裂", "绝缘子缺陷", "塔材腐蚀", "异物悬挂", "导线弧垂异常",
            "管道泄漏", "地面渗透", "植被异常", "设备腐蚀", "阀门损坏",
            "人员定位", "被困人员", "火势蔓延", "危险区域", "道路阻断",
            "水体污染", "植被覆盖", "烟气排放", "垃圾堆积", "野生动物监测"
        ]
        
        missing_targets = [target for target in required_targets if target not in target_name_to_id]
        if missing_targets:
            print(f"错误：以下检测目标不存在，请先运行init_detection_targets.py: {missing_targets}")
            return
        
        # 预设场景数据，使用检测目标ID
        preset_scenarios = [
            {
                "name": "电力巡检",
                "description": "用于高压线路、变电站等电力设施的自动化巡检场景",
                "scenario_type": "preset",
                "drone_type_requirements": [{"name": "多旋翼", "value": "multi_rotor"}, {"name": "固定翼", "value": "fixed_wing"}, {"name": "垂直起降固定翼", "value": "vtol_fixed_wing"}],
                "detection_target_names": ["线路断裂", "绝缘子缺陷", "塔材腐蚀", "异物悬挂", "导线弧垂异常"],
                "default_route_logic": {
                    "type": "waypoint_sequence",
                    "parameters": {
                        "line_following": True,
                        "height_above_line": 50,
                        "side_distance": 20,
                        "photo_interval": 10,
                        "waypoint_spacing": 100
                    }
                },
                "detection_model_ids": [1, 2],
                "permission_config": {
                    "required_roles": ["operator", "inspector"],
                    "min_permission_level": 2
                },
                "is_active": True
            },
            {
                "name": "管道泄漏检测",
                "description": "用于石油、天然气等管道设施的泄漏检测场景",
                "scenario_type": "preset",
                "drone_type_requirements": [{"name": "多旋翼", "value": "multi_rotor"}, {"name": "垂直起降固定翼", "value": "vtol_fixed_wing"}],
                "detection_target_names": ["管道泄漏", "地面渗透", "植被异常", "设备腐蚀", "阀门损坏"],
                "default_route_logic": {
                    "type": "corridor_scan",
                    "parameters": {
                        "corridor_width": 50,
                        "scan_height": 30,
                        "grid_spacing": 20,
                        "thermal_scan": True,
                        "gas_sensor": True
                    }
                },
                "detection_model_ids": [2, 3],
                "permission_config": {
                    "required_roles": ["operator", "inspector", "engineer"],
                    "min_permission_level": 3,
                    "special_permissions": ["hazardous_area_access"]
                },
                "is_active": True
            },
            {
                "name": "应急救援",
                "description": "用于灾害现场、事故救援等应急场景的快速响应",
                "scenario_type": "preset",
                "drone_type_requirements": [{"name": "多旋翼", "value": "multi_rotor"}, {"name": "固定翼", "value": "fixed_wing"}],
                "detection_target_names": ["人员定位", "被困人员", "火势蔓延", "危险区域", "道路阻断"],
                "default_route_logic": {
                    "type": "grid_search",
                    "parameters": {
                        "search_area": {"type": "polygon"},
                        "grid_resolution": 50,
                        "camera_tilt": -45,
                        "photo_quality": "high",
                        "thermal_camera": True,
                        "real_time_transmission": True
                    }
                },
                "detection_model_ids": [1, 3],
                "permission_config": {
                    "required_roles": ["rescuer", "commander", "operator"],
                    "min_permission_level": 2,
                    "emergency_override": True
                },
                "is_active": True
            },
            {
                "name": "环保监测",
                "description": "用于环境监测、水质分析、空气质量检测等环保场景",
                "scenario_type": "preset",
                "drone_type_requirements": [{"name": "多旋翼", "value": "multi_rotor"}, {"name": "垂直起降固定翼", "value": "vtol_fixed_wing"}],
                "detection_target_names": ["水体污染", "植被覆盖", "烟气排放", "垃圾堆积", "野生动物监测"],
                "default_route_logic": {
                    "type": "area_coverage",
                    "parameters": {
                        "waypoints": [],
                        "altitude": 100,
                        "camera_settings": {
                            "resolution": "high",
                            "NDVI_filter": True
                        },
                        "sensor_reading_interval": 5,
                        "sample_collection": True
                    }
                },
                "detection_model_ids": [2],
                "permission_config": {
                    "required_roles": ["scientist", "operator", "analyst"],
                    "min_permission_level": 1
                },
                "is_active": True
            }
        ]
        
        if is_update_mode:
            # 更新现有预设场景
            updated_count = 0
            for scenario in existing_presets:
                # 查找匹配的预设数据
                for scenario_data in preset_scenarios:
                    if scenario.name == scenario_data["name"]:
                        # 将检测目标名称转换为ID
                        detection_target_ids = [target_name_to_id[name] for name in scenario_data["detection_target_names"]]
                        
                        # 先删除现有场景，然后创建新的（确保数据格式正确）
                        db.delete(scenario)
                        
                        # 创建新的场景
                        new_scenario = TaskScenario(
                            name=scenario_data["name"],
                            description=scenario_data["description"],
                            scenario_type=scenario_data["scenario_type"],
                            drone_type_requirements=json.dumps(scenario_data["drone_type_requirements"]),
                            detection_targets=None,  # 不再使用JSON格式存储检测目标
                            default_route_logic=json.dumps(scenario_data["default_route_logic"]),
                            detection_model_ids=json.dumps(scenario_data["detection_model_ids"]),
                            permission_config=json.dumps(scenario_data["permission_config"]),
                            is_active=scenario_data["is_active"]
                        )
                        
                        # 关联检测目标
                        targets = db.query(DetectionTarget).filter(DetectionTarget.id.in_(detection_target_ids)).all()
                        new_scenario.detection_target_objects = targets
                        
                        db.add(new_scenario)
                        updated_count += 1
                        break
                        
                        # 检测目标关联已在创建新场景时处理
                        break
            
            # 提交事务
            db.commit()
            print(f"成功更新{updated_count}个预设场景模板")
        else:
            # 创建新的预设场景
            for scenario_data in preset_scenarios:
                # 将检测目标名称转换为ID
                detection_target_ids = [target_name_to_id[name] for name in scenario_data["detection_target_names"]]
                
                # 创建场景
                scenario = TaskScenario(
                    name=scenario_data["name"],
                    description=scenario_data["description"],
                    scenario_type=scenario_data["scenario_type"],
                    drone_type_requirements=json.dumps(scenario_data["drone_type_requirements"]),
                    detection_targets=None,  # 不再使用JSON格式存储检测目标
                    default_route_logic=json.dumps(scenario_data["default_route_logic"]),
                    detection_model_ids=json.dumps(scenario_data["detection_model_ids"]),
                    permission_config=json.dumps(scenario_data["permission_config"]),
                    is_active=scenario_data["is_active"]
                )
                
                # 关联检测目标
                targets = db.query(DetectionTarget).filter(DetectionTarget.id.in_(detection_target_ids)).all()
                scenario.detection_target_objects = targets
                
                db.add(scenario)
            
            # 提交事务
            db.commit()
            print(f"成功初始化{len(preset_scenarios)}个预设场景模板")
        
    except Exception as e:
        print(f"初始化任务场景数据失败: {e}")
        db.rollback()
        raise
    finally:
        db.close()

if __name__ == "__main__":
    init_task_scenarios()